import os
import logging
from concurrent.futures import Future, ThreadPoolExecutor

from tinyops.config import DirectoryManager
from tinyops.library.executor import SSHExecutor, ContainerExecutor, LocalExecutor
from tinyops.library.task import (
    TaskEnv,
    TargetType,
    TaskCommand,
    TaskStep,
    TaskTarget,
    TaskLogger,
)

from tinyops.config import Config
from tinyops.notifier import NotifierManager

from .task import Task


LOG = logging.getLogger(__name__)


class TaskExecutor:
    def __init__(self, task: Task, logger: TaskLogger = None):
        self.task = task
        self.config: Config = task.config
        dm = DirectoryManager(self.config)
        self.work_dir = dm.task_dir(task.name)
        notifier_factory = NotifierManager(task.config)
        self.notifier = notifier_factory.create(*task.notifier)
        if logger is None:
            logger = TaskLogger()
        self.logger = logger
        self.task_env = TaskEnv(self.task.env)
        self.targets = {}
        self.thread_pool = ThreadPoolExecutor(3)

    def _worker(self, target, step: TaskStep) -> bool:
        target.init()
        handlers = {
            TaskCommand.RUN: getattr(target, 'do_run', None),
            TaskCommand.GET: getattr(target, 'do_get', None),
            TaskCommand.PUT: getattr(target, 'do_put', None),
        }
        handler = handlers.get(step.command)
        if not handler:
            self.logger.log(
                target.name, f'unknown command {step.command}', sep='>')
            return
        return handler(**step.params)

    def _resolve_targets(self, step):
        server_targets = {}
        docker_targets = {}
        local_target = None
        target_info: TaskTarget
        for target_info in step.targets:
            if target_info.type == TargetType.LOCAL:
                local_target = True
            if target_info.type == TargetType.SERVER:
                for x in self.config.server(target_info.name):
                    server_targets[x.name] = x
            if target_info.type == TargetType.DOCKER:
                for x in self.config.docker(target_info.name):
                    docker_targets[(x.name, target_info.image)] = x
        common_kwargs = dict(
            env=self.task_env,
            timeout=self.task.timeout,
            logger=self.logger,
        )
        targets = []
        if local_target:
            name = 'local'
            target = self.targets.get(name)
            if not target:
                target = LocalExecutor(**common_kwargs)
                self.targets[name] = target
            targets.append(target)
        for name, x in server_targets.items():
            target = self.targets.get(name)
            if not target:
                target = SSHExecutor(x, **common_kwargs)
                self.targets[name] = target
            targets.append(target)
        for (name, image), x in docker_targets.items():
            target = self.targets.get((name, image))
            if not target:
                target = ContainerExecutor(x, image=image, **common_kwargs)
                self.targets[(name, image)] = target
            targets.append(target)
        return targets

    def _execute_step(self, step: TaskStep) -> bool:
        if step.command == TaskCommand.LOG:
            return self.do_log(**step.params)
        if step.command == TaskCommand.NOTIFY:
            return self.do_notify(**step.params)
        futures = []
        targets = self._resolve_targets(step)
        fut: Future
        for target in targets:
            fut = self.thread_pool.submit(self._worker, target, step)
            futures.append(fut)
        ok_s = [fut.result() for fut in futures]
        return all(ok_s)

    def do_log(self, message):
        self.logger.log('*', f'## {message}', sep='>')
        return True

    def do_notify(self, message):
        self.logger.log('*', f'## NOTIFY {message}', sep='>')
        self.notifier.notify(message)
        return True

    def _close(self):
        self.thread_pool.shutdown(wait=False)
        for x in self.targets.values():
            x.close()

    def execute(self) -> bool:
        current_work_dir = os.getcwd()
        is_success = True
        try:
            os.chdir(self.work_dir)
            for step in self.task.steps:
                is_ok = self._execute_step(step)
                if not is_ok:
                    is_success = False
                    msg = f'Failed at step line#{step.lineno}'
                    self.logger.log('*', msg, sep='*')
                    break
                self.task_env.apply_export()
        finally:
            os.chdir(current_work_dir)
            self._close()
        return is_success
